#pragma once

#include <string>
#include <map>
#include <mutex>
#include <optional>
#include <filesystem>
#include <variant>

#include <Poco/Timestamp.h>
#include <Poco/Util/AbstractConfiguration.h>
#include <Core/Defines.h>
#include <IO/ReadSettings.h>
#include <IO/WriteSettings.h>
#include <IO/copyData.h>

#include <Core/Types.h>
#include <Common/Exception.h>
#include <Common/ObjectStorageKey.h>
#include <Common/ObjectStorageKeyGenerator.h>
#include <Common/ThreadPool.h>
#include <Common/ThreadPool_fwd.h>
#include <Common/threadPoolCallbackRunner.h>

#include <Disks/DirectoryIterator.h>
#include <Disks/DiskType.h>
#include <Disks/DiskObjectStorage/ObjectStorages/StoredObject.h>
#include <Disks/WriteMode.h>

#include <Processors/ISimpleTransform.h>
#include <Processors/Formats/IInputFormat.h>
#include <Storages/ObjectStorage/DataLakes/DataLakeObjectMetadata.h>

#include <Interpreters/Context_fwd.h>
#include "config.h"

#if USE_AZURE_BLOB_STORAGE

#include <azure/core/credentials/credentials.hpp>
#include <azure/storage/common/storage_credential.hpp>
#include <azure/identity/managed_identity_credential.hpp>
#include <azure/identity/workload_identity_credential.hpp>
#include <azure/identity/client_secret_credential.hpp>

namespace DB::AzureBlobStorage
{

class ContainerClientWrapper;
using ContainerClient = ContainerClientWrapper;

class StaticCredential : public Azure::Core::Credentials::TokenCredential
{
public:
    StaticCredential(std::string token_, std::chrono::system_clock::time_point expires_on_)
        : token(std::move(token_)), expires_on(expires_on_)
    {}

    Azure::Core::Credentials::AccessToken GetToken(
        Azure::Core::Credentials::TokenRequestContext const &,
        Azure::Core::Context const &) const override
    {
        return Azure::Core::Credentials::AccessToken { .Token = token, .ExpiresOn = expires_on };
    }

private:
    std::string token;
    std::chrono::system_clock::time_point expires_on;
};

using ConnectionString = StrongTypedef<String, struct ConnectionStringTag>;

using AuthMethod = std::variant<
    ConnectionString,
    std::shared_ptr<Azure::Identity::ClientSecretCredential>,
    std::shared_ptr<Azure::Storage::StorageSharedKeyCredential>,
    std::shared_ptr<Azure::Identity::WorkloadIdentityCredential>,
    std::shared_ptr<Azure::Identity::ManagedIdentityCredential>,
    std::shared_ptr<AzureBlobStorage::StaticCredential>>;
}

#endif

#if USE_AWS_S3
namespace DB::S3
{
class Client;
}
#endif

namespace DB
{

namespace ErrorCodes
{
extern const int NOT_IMPLEMENTED;
}

class ReadBufferFromFileBase;
class WriteBufferFromFileBase;

using ObjectAttributes = std::map<std::string, std::string>;

struct ObjectMetadata
{
    uint64_t size_bytes = 0;
    Poco::Timestamp last_modified;
    std::string etag;
    ObjectAttributes tags;
    ObjectAttributes attributes;
};

struct DataLakeObjectMetadata;

struct RelativePathWithMetadata
{
    String relative_path;
    /// Object metadata: size, modification time, etc.
    std::optional<ObjectMetadata> metadata;

    RelativePathWithMetadata() = default;

    explicit RelativePathWithMetadata(String relative_path_, std::optional<ObjectMetadata> metadata_ = std::nullopt)
        : relative_path(std::move(relative_path_))
        , metadata(std::move(metadata_))
    {}

    RelativePathWithMetadata(const RelativePathWithMetadata & other) = default;

    ~RelativePathWithMetadata() = default;

    std::string getFileName() const { return std::filesystem::path(relative_path).filename(); }
    std::string getPath() const { return relative_path; }
};

struct ObjectKeyWithMetadata
{
    ObjectStorageKey key;
    ObjectMetadata metadata;

    ObjectKeyWithMetadata() = default;

    ObjectKeyWithMetadata(ObjectStorageKey key_, ObjectMetadata metadata_)
        : key(std::move(key_))
        , metadata(std::move(metadata_))
    {}
};

struct SmallObjectDataWithMetadata
{
    std::string data;
    ObjectMetadata metadata;
};

using RelativePathWithMetadataPtr = std::shared_ptr<RelativePathWithMetadata>;
using RelativePathsWithMetadata = std::vector<RelativePathWithMetadataPtr>;
using ObjectKeysWithMetadata = std::vector<ObjectKeyWithMetadata>;

class IObjectStorageIterator;
using ObjectStorageIteratorPtr = std::shared_ptr<IObjectStorageIterator>;

/// Base class for all object storages which implement some subset of ordinary filesystem operations.
///
/// Examples of object storages are S3, Azure Blob Storage, HDFS.
class IObjectStorage
{
public:
    IObjectStorage() = default;

    virtual std::string getName() const = 0;

    virtual ObjectStorageType getType() const = 0;

    /// The logical root or base path used to group a set of related objects.
    virtual std::string getRootPrefix() const { return ""; }

    /// Common object key prefix relative to the root path.
    virtual std::string getCommonKeyPrefix() const = 0;

    virtual std::string getDescription() const = 0;

    /// Object exists or not
    virtual bool exists(const StoredObject & object) const = 0;

    /// Object exists or any child on the specified path exists.
    /// We have this method because object storages are flat for example
    /// /a/b/c/d may exist but /a/b/c may not. So this method will return true for
    /// /, /a, /a/b, /a/b/c, /a/b/c/d while exists will return true only for /a/b/c/d
    virtual bool existsOrHasAnyChild(const std::string & path) const;

    /// List objects recursively by certain prefix.
    virtual void listObjects(const std::string & path, RelativePathsWithMetadata & children, size_t max_keys) const;

    /// List objects recursively by certain prefix. Use it instead of listObjects, if you want to list objects lazily.
    virtual ObjectStorageIteratorPtr iterate(const std::string & path_prefix, size_t max_keys, bool with_tags) const;

    /// Get object metadata if supported. It should be possible to receive at least size of object
    virtual ObjectMetadata getObjectMetadata(const std::string & path, bool with_tags) const = 0;

    /// Same as getObjectMetadata(), but ignores if object does not exist.
    virtual std::optional<ObjectMetadata> tryGetObjectMetadata(const std::string & path, bool with_tags) const = 0;

    /// Read single object
    virtual std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
        const StoredObject & object,
        const ReadSettings & read_settings,
        std::optional<size_t> read_hint = {}) const = 0;

    /// Read small object into memory and return it as string
    /// Also contain consistent object metadata if available in this object storage.
    /// if size of object is larger than max_size_bytes throws exception
    ///
    /// NOTE: This method exists because it's impossible to get object metadata in generic way
    /// from readObject. ReadObject returns ReadBufferFromFileBase and most of implementations
    /// issue first request to object store only in first call of read/next method.
    ///
    /// WARN: Don't use this method for large objects, it will eat all memory.
    /// That is the reason why max_size_bytes is explicit and 0-value will not help to bypass this check.
    virtual SmallObjectDataWithMetadata readSmallObjectAndGetObjectMetadata( /// NOLINT
        const StoredObject & object,
        const ReadSettings & read_settings,
        size_t max_size_bytes,
        std::optional<size_t> read_hint = {}) const;

    /// Open the file for write and return WriteBufferFromFileBase object.
    virtual std::unique_ptr<WriteBufferFromFileBase> writeObject( /// NOLINT
        const StoredObject & object,
        WriteMode mode,
        std::optional<ObjectAttributes> attributes = {},
        size_t buf_size = DBMS_DEFAULT_BUFFER_SIZE,
        const WriteSettings & write_settings = {}) = 0;

    virtual bool isRemote() const = 0;

    /// Remove object on path if exists
    virtual void removeObjectIfExists(const StoredObject & object) = 0;

    /// Remove objects on path if exists
    virtual void removeObjectsIfExist(const StoredObjects & object) = 0;

    /// Copy object with different attributes if required
    virtual void copyObject( /// NOLINT
        const StoredObject & object_from,
        const StoredObject & object_to,
        const ReadSettings & read_settings,
        const WriteSettings & write_settings,
        std::optional<ObjectAttributes> object_to_attributes = {}) = 0;

    /// Copy object to another instance of object storage
    /// by default just read the object from source object storage and write
    /// to destination through buffers.
    virtual void copyObjectToAnotherObjectStorage( /// NOLINT
        const StoredObject & object_from,
        const StoredObject & object_to,
        const ReadSettings & read_settings,
        const WriteSettings & write_settings,
        IObjectStorage & object_storage_to,
        std::optional<ObjectAttributes> object_to_attributes = {});

    virtual ~IObjectStorage() = default;

    virtual const std::string & getCacheName() const;

    static ThreadPool & getThreadPoolWriter();

    virtual void shutdown() = 0;

    virtual void startup() = 0;

    /// Apply new settings, in most cases reiniatilize client and some other staff
    struct ApplyNewSettingsOptions
    {
        bool allow_client_change = true;
    };
    virtual void applyNewSettings(
        const Poco::Util::AbstractConfiguration & /* config */,
        const std::string & /*config_prefix*/,
        ContextPtr /* context */,
        const ApplyNewSettingsOptions & /* options */) {}

    /// Sometimes object storages have something similar to chroot or namespace, for example
    /// buckets in S3. If object storage doesn't have any namepaces return empty string.
    virtual String getObjectsNamespace() const = 0;

    /// Get unique id for passed absolute path in object storage.
    virtual std::string getUniqueId(const std::string & path) const { return path; }

    /// Remove filesystem cache.
    virtual void removeCacheIfExists(const std::string & /* path */) {}

    virtual bool supportsCache() const { return false; }

    virtual bool isReadOnly() const { return false; }

    virtual bool supportParallelWrite() const { return false; }

    virtual ReadSettings patchSettings(const ReadSettings & read_settings) const;

    virtual WriteSettings patchSettings(const WriteSettings & write_settings) const;

    virtual ObjectStorageKeyGeneratorPtr createKeyGenerator() const = 0;

#if USE_AZURE_BLOB_STORAGE
    virtual std::shared_ptr<const AzureBlobStorage::ContainerClient> getAzureBlobStorageClient() const
    {
        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "This function is only implemented for AzureBlobStorage");
    }

    virtual AzureBlobStorage::AuthMethod getAzureBlobStorageAuthMethod() const
    {
        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "This function is only implemented for AzureBlobStorage");
    }
#endif

#if USE_AWS_S3
    virtual std::shared_ptr<const S3::Client> getS3StorageClient()
    {
        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "This function is only implemented for S3ObjectStorage");
    }
    virtual std::shared_ptr<const S3::Client> tryGetS3StorageClient() { return nullptr; }
#endif

#if USE_AZURE_BLOB_STORAGE || USE_AWS_S3
    /// Assign tag on objects
    virtual void tagObjects(const StoredObjects &, const std::string &, const std::string &)
    {
        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "The method 'tagObjects' is only implemented for S3 and Azure storages");
    }
#endif
};

using ObjectStoragePtr = std::shared_ptr<IObjectStorage>;

}
